import math
import arrow

from app import celery_app
from app.db.models.tasks import Task
from app.db.models.jobs import Job
from app.db.models.asset_tasks import AssetTasks
from app.db.models.tasks import IgnoreWarning
from app.db.models.users import User
from app.db.models.wass import ScanEvents
from app.api.v2.job_util import update_last_web_security_status
from app.libs.wechat import wecom_notify
from app.libs.display import SEVERITY_LEVEL_MAP, IGNORE_RISK_REASON
from app.libs.enums import TASK_TYPE_CONFIG, TaskType
from app.config.settings import DATETIME_FMT


def task_warning_counts(warnings):
    res = {
        "warningCount": len([w for w in warnings if not w.ignoredTime]),
        "warningResolvedCount": len([w for w in warnings if not w.ignoredTime and w.isResolved])
    }
    res["warningUnresolvedCount"] = res["warningCount"] - res["warningResolvedCount"]
    return res


@celery_app.task(acks_late=True, name="scanv:after_ignore_task_warning")
def after_ignore_task_warning(task_id):
    """忽略威胁后重新计算资产、任务组等数据"""
    task_obj = Task.objects.filter(id=task_id, status="completed").first()
    if not task_obj:
        return
    warnings = task_obj.result.warnings
    if (getattr(task_obj.result, "targetStatus", None) and
            task_obj.result.targetStatus.status != "good" and len(warnings) == 0):
        securityStatus = "unknown"
    else:
        securityStatus = "warning" if len(warnings) > 0 else "safe"
    task_update_dict = {"securityStatus": securityStatus}
    for key, value in task_warning_counts(warnings).items():
        task_update_dict.update({f"result.{key}": value})
    Task._get_collection().update_one(
        {"_id": task_obj.id},
        {"$set": task_update_dict}
    )
    task_type = task_obj.taskType
    query = {f"{task_type}ResultId": task_obj.id, "uid": task_obj.uid}
    if job_obj := Job.objects.find(query).first():
        update_last_web_security_status(job_obj)
    if asset_task_obj := AssetTasks.objects.find(query).first():
        asset_task_obj._get_web_security_last_result(is_save=True)


@celery_app.task(acks_late=True, name="scanv:notify_wecom_ignore_warning")
def notify_wecom_ignore_warning(uid, task_session_id, warning_ids):
    data = []
    user = User.objects.filter(id=uid).first()
    if not user:
        return
    task_ids = IgnoreWarning.objects.filter(uid=uid, detail__id__in=warning_ids).distinct("taskId")
    if not task_ids:
        return
    scan_event_db = ScanEvents._get_collection()
    for w in IgnoreWarning.objects.filter(uid=uid, detail__id__in=warning_ids):
        item = [
            f"用户名: {user.username}",
            f"资产地址: {w.targetUrl}",
            f"任务类型: {TASK_TYPE_CONFIG.get(w.taskType, {}).get('cn')}",
            f"忽略原因: {IGNORE_RISK_REASON.get(w.reason.type)}",
            f"忽略结果ID: {str(w.id)}",
            f"忽略时间: {arrow.get(w.ignoredTime).to('local').datetime.strftime(DATETIME_FMT)}",
            f"威胁地址: {w.detail.affects}",
            f"威胁名称: {w.detail.title}",
            f"威胁等级: {SEVERITY_LEVEL_MAP[w.detail.level] if w.detail.level >= 2 else '低危'}",
            f"task_session_id: {task_session_id}",
            f"event_id: "
        ]
        base_q = {"task_session_id": task_session_id}
        q = {}
        if w.taskType == TaskType.ssl.value:
            q = {"detail.name.zh_cn": w.detail.title, **base_q}
        elif w.taskType == TaskType.vul.value:
            q = {
                "detail.name.zh_cn": w.detail.title,
                "$or": [{"detail.affect": w.detail.affects}, {"detail.url": w.detail.affects}],
                **base_q
            }
        elif w.taskType == TaskType.change_check.value:
            q = {"detail.key": "change_check", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "foreign_links":
            q = {"event_name": "risk_link", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "broken_links":
            q = {"event_name": "statistics", "detail.broken_links.referer": w.detail.affects, **base_q}
        elif w.detail.category == "black_links":
            q = {"event_name": "black_links", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "cryjack":
            q = {"event_name": "cryptojacking", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "malscan":
            q = {"event_name": "malscan", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "keyword":
            q = {"event_name": "keyword", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "privacy_disclosure":
            q = {"event_name": "privacy_disclosure", "detail.url": w.detail.affects, **base_q}
        if q and (event_ids := [e["event_id"] for e in scan_event_db.find(q)]):
            item[-1] = f"event_id: {'、'.join(event_ids)}"
        data.append("\n".join(item))

    if not data:
        return
    # 每次发送五条，多了可能微信单条消息字数限制
    count = len(data)
    pages = math.ceil(count / 5)
    for i in range(0, pages):
        wecom_notify("\n\n".join(data[i * 5: (i+1) * 5]))
